import os
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel, QPushButton, QHBoxLayout, QLineEdit
from PyQt5.QtGui import QPixmap, QFont
from PyQt5.QtGui import QImage, QColor
from PyQt5.QtCore import Qt
from PyQt5 import QtCore

class SameNameImageViewer(QWidget):
    size=400
    def __init__(self, parent_folder):
        super().__init__()
        self.parent_folder = parent_folder
        self.image_groups = {}
        self.current_group_index = 0
        self.initUI()

    def initUI(self):
        self.find_same_named_images()
        self.display_images()
        self.create_buttons()
        self.create_image_input()
        self.setMinimumSize(400, 300)
        self.setWindowTitle('同名图片查看器')
        self.setStyleSheet("background-color: white;")
        self.show()

    def create_image_input(self):
        self.input_layout = QHBoxLayout()
        self.image_input = QLineEdit()
        self.input_layout.addWidget(self.image_input)
        self.show_button = QPushButton("显示图片")
        self.show_button.clicked.connect(self.show_image_by_name)
        self.input_layout.addWidget(self.show_button)
        self.layout.addLayout(self.input_layout, len(self.image_groups) * 2 + 2, 0, 1, 4)

    def show_image_by_name(self):
        image_name = self.image_input.text()
        print(image_name)
        if image_name in self.image_groups:
            if self.layout is not None:
                for i in reversed(range(self.layout.count())):
                    widget = self.layout.itemAt(i).widget()
                    if widget is not None:
                        widget.setParent(None)
        row, col = 0, 0
        group_name = image_name
        group_label = QLabel(f"图: {group_name}")
        group_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
        self.layout.addWidget(group_label, row, col, 1, 4)
        row += 1
        col = 0
        for image_info in self.image_groups[group_name]:
            image_path, parent_folder = image_info
            image_label = QLabel()
            pixmap = QPixmap(image_path)
            pixmap = pixmap.scaled(256, 256)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移

            # 获取 val_label_demo 文件夹对应的图像路径
            a_folder_path = os.path.join(self.parent_folder, 'label/predictResult')
            a_image_path = os.path.join(a_folder_path, os.path.basename(image_path))

            # 设置不需要对比的图片
            if parent_folder in ['image', 'label']:
                tmp = pixmap.scaled(self.size, self.size)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移
                image_label.setPixmap(tmp)
            else:
                # 加载图像并进行像素级别的比较
                result_pixmap = self.compare_images(pixmap, QPixmap(a_image_path))
                result_pixmap = result_pixmap.scaled(self.size, self.size)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移
                image_label.setPixmap(result_pixmap)

            # image_label.setScaledContents(True)  # 将 pixmap 缩放以适应 QLabel
            self.layout.addWidget(image_label, row, col)
            font = QFont()
            font.setPointSize(8)
            folder_label = QLabel(parent_folder)
            folder_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
            folder_label.setFont(font)  # 设置字体大小
            self.layout.addWidget(folder_label, row + 1, col)
            col = (col + 1) % 4
            if col == 0:
                row += 2

        self.setLayout(self.layout)

    def find_same_named_images(self):

        # list
        folder_names = [
            'image',
            'label',
            'SegNet',
            'U_Net',
            'PSPNet',
            'ShiftPoolingPSPNet',
            'NestedUNet',
            'Nswintransformer',

        ]

        # 遍历列表中的文件夹名称，依次处理
        for folder_name in folder_names:
            folder_path = os.path.join(self.parent_folder, folder_name, 'predictResult')
            # 检查文件夹是否存在
            if not os.path.exists(folder_path):
                continue
            # 调用处理文件夹的方法
            self.process_folder(folder_path)


    def process_folder(self, folder_path):
        for file in os.listdir(folder_path):
            if file.endswith('.png') or file.endswith('.jpg'):
                image_path = os.path.join(folder_path, file)
                image_name = os.path.splitext(file)[0]
                parent_folder = os.path.basename(os.path.dirname(folder_path))
                if image_name not in self.image_groups:
                    self.image_groups[image_name] = []
                self.image_groups[image_name].append((image_path, parent_folder))

    def display_images(self):
        self.layout = QGridLayout(self)
        self.update_display()

    def update_display(self):
        if self.layout is not None:
            for i in reversed(range(self.layout.count())):
                widget = self.layout.itemAt(i).widget()
                if widget is not None:
                    widget.setParent(None)

        row, col = 0, 0
        group_name = list(self.image_groups.keys())[self.current_group_index]
        group_label = QLabel(f"图: {group_name}")
        group_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
        self.layout.addWidget(group_label, row, col, 1, 4)
        row += 1
        col = 0
        for image_info in self.image_groups[group_name]:
            image_path, parent_folder = image_info
            image_label = QLabel()
            pixmap = QPixmap(image_path)
            pixmap = pixmap.scaled(256,256)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移

            # 获取 val_label_demo 文件夹对应的图像路径
            a_folder_path = os.path.join(self.parent_folder, 'label/predictResult')
            a_image_path = os.path.join(a_folder_path, os.path.basename(image_path))

            #设置不需要对比的图片
            if parent_folder in ['image', 'label']:
                tmp=pixmap.scaled(self.size,self.size)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移
                image_label.setPixmap(tmp)
            else:
                # 加载图像并进行像素级别的比较
                result_pixmap = self.compare_images(pixmap, QPixmap(a_image_path))
                result_pixmap=result_pixmap.scaled(self.size,self.size)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移
                image_label.setPixmap(result_pixmap)



            # image_label.setScaledContents(True)  # 将 pixmap 缩放以适应 QLabel
            self.layout.addWidget(image_label, row, col)
            font = QFont()
            font.setPointSize(8)
            folder_label = QLabel(parent_folder)
            folder_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
            folder_label.setFont(font)  # 设置字体大小
            self.layout.addWidget(folder_label, row + 1, col)
            col = (col + 1) % 4
            if col == 0:
                row += 2

        self.setLayout(self.layout)

    def compare_images(self, pixmap1, pixmap2):
        # 将 QPixmap 转换为 QImage
        image1 = pixmap1.toImage()
        image2 = pixmap2.toImage()
        width = image1.width()
        height = image1.height()
        result_image = QImage(width, height, QImage.Format_ARGB32)

        # 遍历像素进行比较
        for y in range(height):
            for x in range(width):
                pixel1 = QColor(image1.pixel(x, y))
                pixel2 = QColor(image2.pixel(x, y))


                # 如果像素不同，根据像素的颜色分别判断 TP、FP、TN、FN，并设置相应的颜色
                if pixel1 == QColor(Qt.white) and pixel2 == QColor(Qt.black):
                    result_image.setPixelColor(x, y, QColor(Qt.green))  # FN
                elif pixel1 == QColor(Qt.black) and pixel2 == QColor(Qt.white):
                    result_image.setPixelColor(x, y, QColor(Qt.red))  # FP
                elif pixel1 == QColor(Qt.black) and pixel2 == QColor(Qt.black):
                    result_image.setPixelColor(x, y, QColor(Qt.black))  # TN
                elif pixel1 == QColor(Qt.white) and pixel2 == QColor(Qt.white):
                    result_image.setPixelColor(x, y, QColor(Qt.white))  # TP

        return QPixmap.fromImage(result_image)

    def create_buttons(self):
        button_layout = QHBoxLayout()

        prev_button = QPushButton("上一组")
        prev_button.clicked.connect(self.previous_group)
        button_layout.addWidget(prev_button)

        next_button = QPushButton("下一组")
        next_button.clicked.connect(self.next_group)
        button_layout.addWidget(next_button)

        self.layout.addLayout(button_layout, len(self.image_groups) * 2 + 1, 0, 1, 4)

    def previous_group(self):
        if self.current_group_index > 0:
            self.current_group_index -= 1
            self.update_display()

    def next_group(self):
        if self.current_group_index < len(self.image_groups) - 1:
            self.current_group_index += 1
            self.update_display()

if __name__ == '__main__':
    app = QApplication([])
    parent_folder = '/media/jolly/HIK/THE FLAME DATASET'  # 替换为您的父文件夹路径
    window = SameNameImageViewer(parent_folder)
    app.exec_()
